import argparse
import os
import re
import sys
import glob
from pathlib import Path
import numpy as np
import torch
import time
from torch import nn 
from eval_utils import eval_utils
from pcdet.config import cfg, cfg_from_list, cfg_from_yaml_file, update_cfg_by_args, log_config_to_file
from pcdet.datasets import build_dataloader
from pcdet.models import build_network
from pcdet.utils import common_utils
from pcdet.utils import calibration_kitti as calibration

import typing
from torch.onnx import symbolic_helper
from torch.onnx.symbolic_helper import parse_args

torch.backends.cudnn.benchmark = True


_OPSET_VERSION = 11
_registered_ops: typing.AbstractSet[str] = set()

def parse_config():
    parser = argparse.ArgumentParser(description='arg parser')
    parser.add_argument('--cfg_file', type=str, default=None, help='specify the config for training')

    # basic testing options
    parser.add_argument('--batch_size', type=int, default=None, required=False, help='batch size for training')
    parser.add_argument('--workers', type=int, default=2, help='number of workers for dataloader')
    parser.add_argument('--exp_name', type=str, default='default', help='exp path for this experiment')
    parser.add_argument('--eval_tag', type=str, default='default', help='eval tag for this experiment')
    parser.add_argument('--max_waiting_mins', type=int, default=30, help='max waiting minutes')
    parser.add_argument('--save_to_file', action='store_true', default=False, help='')
    # loading options
    parser.add_argument('--ckpt', type=str, default=None, help='checkpoint to evaluate')
    parser.add_argument('--load_eval', action='store_true', default=False, help='')
    parser.add_argument('--ckpt_id', type=int, default=None, help='checkpoint id to evaluate')
    parser.add_argument('--start_epoch', type=int, default=0, help='')
    # distributed options
    parser.add_argument('--launcher', choices=['none', 'pytorch', 'slurm'], default='none')
    parser.add_argument('--tcp_port', type=int, default=18888, help='tcp port for distrbuted training')
    parser.add_argument('--local_rank', type=int, default=0, help='local rank for distributed training')
    # config options
    parser.add_argument('--set', dest='set_cfgs', default=None, nargs=argparse.REMAINDER, help='set extra config keys if needed')
    parser.add_argument('--trainval', action='store_true', default=False, help='')
    parser.add_argument('--imitation', type=str, default="2d")
    parser.add_argument('--save_depth', action='store_true', default=False)
    parser.add_argument('--convert_onnx', action='store_true', default=False)

    args = parser.parse_args()

    cfg_from_yaml_file(args.cfg_file, cfg)
    update_cfg_by_args(cfg, args)
    cfg.TAG = Path(args.cfg_file).stem
    cfg.EXP_GROUP_PATH = '_'.join(args.cfg_file.split('/')[1:-1])  # remove 'cfgs' and 'xxxx.yaml'

    np.random.seed(1024)

    assert args.ckpt or args.ckpt_id, "pls specify ckpt or ckpt_dir or ckpt_id"

    if args.set_cfgs is not None:
        cfg_from_list(args.set_cfgs, cfg)

    return args, cfg


def eval_single_ckpt(model, test_loader, args, eval_output_dir, logger, epoch_id, dist_test=False):
    # load checkpoint
    model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=dist_test)
    model.cuda()
    if args.load_eval:
        eval_utils.eval_one_epoch_load(
            cfg, model, test_loader, epoch_id, logger, dist_test=dist_test,
            result_dir=eval_output_dir, save_to_file=args.save_to_file
        )
    else:
        # import pdb
        # pdb.set_trace()
        eval_utils.eval_one_epoch(
            cfg, model, test_loader, epoch_id, logger, dist_test=dist_test,
            result_dir=eval_output_dir, save_to_file=args.save_to_file,
            save_depth=args.save_depth, convert_onnx=args.convert_onnx
        )

def get_no_evaluated_ckpt(ckpt_dir, ckpt_record_file, args):
    ckpt_list = glob.glob(os.path.join(ckpt_dir, '*checkpoint_epoch_*.pth'))
    ckpt_list.sort(key=os.path.getmtime)
    evaluated_ckpt_list = [float(x.strip()) for x in open(ckpt_record_file, 'r').readlines()]

    for cur_ckpt in ckpt_list:
        num_list = re.findall('checkpoint_epoch_(.*).pth', cur_ckpt)
        if num_list.__len__() == 0:
            continue

        epoch_id = num_list[-1]
        if 'optim' in epoch_id:
            continue
        # if float(epoch_id) not in evaluated_ckpt_list and int(float(epoch_id)) >= args.start_epoch:
        if float(epoch_id) not in evaluated_ckpt_list and int(float(epoch_id)) >= args.start_epoch and int(float(epoch_id)) % args.eval_interval == 0:
            return epoch_id, cur_ckpt
    return -1, None


def repeat_eval_ckpt(model, test_loader, args, eval_output_dir, logger, ckpt_dir, dist_test=False):
    # evaluated ckpt record
    ckpt_record_file = eval_output_dir / ('eval_list_%s.txt' % cfg.DATA_CONFIG.DATA_SPLIT['test'])
    with open(ckpt_record_file, 'a'):
        pass

    # tensorboard log
    # if cfg.LOCAL_RANK == 0:
    #     tb_log = SummaryWriter(log_dir=str(eval_output_dir / ('tensorboard_%s' % cfg.DATA_CONFIG.DATA_SPLIT['test'])))
    total_time = 0
    first_eval = True

    while True:
        # check whether there is checkpoint which is not evaluated
        cur_epoch_id, cur_ckpt = get_no_evaluated_ckpt(ckpt_dir, ckpt_record_file, args)
        if cur_epoch_id == -1 or int(float(cur_epoch_id)) < args.start_epoch:
            wait_second = 30
            if cfg.LOCAL_RANK == 0:
                print('Wait %s seconds for next check (progress: %.1f / %d minutes): %s \r'
                      % (wait_second, total_time * 1.0 / 60, args.max_waiting_mins, ckpt_dir), end='', flush=True)
            time.sleep(wait_second)
            total_time += 30
            if total_time > args.max_waiting_mins * 60 and (first_eval is False):
                break
            continue

        total_time = 0
        first_eval = False

        model.load_params_from_file(filename=cur_ckpt, logger=logger, to_cpu=dist_test)
        model.cuda()

        # start evaluation
        cur_result_dir = eval_output_dir / ('epoch_%s' % cur_epoch_id) / cfg.DATA_CONFIG.DATA_SPLIT['test']
        # tb_dict =
        eval_utils.eval_one_epoch(
            cfg, model, test_loader, cur_epoch_id, logger, dist_test=dist_test,
            result_dir=cur_result_dir, save_to_file=args.save_to_file
        )

        # if cfg.LOCAL_RANK == 0:
        #     for key, val in tb_dict.items():
        #         tb_log.add_scalar(key, val, cur_epoch_id)

        # record this epoch which has been evaluated
        with open(ckpt_record_file, 'a') as f:
            print('%s' % cur_epoch_id, file=f)
        logger.info('Epoch %s has been evaluated' % cur_epoch_id)


class StereoDSGN_post(nn.Module):
    def __init__(self,model):
        super(StereoDSGN_post,self).__init__()
        self.model = model

    def forward(self, left_img, right_img, fu_mul_baseline, calibs_Proj, norm_coord_img, valid_mask_2d, valids):
        batch_dict = {
            'left_img': torch.as_tensor(left_img, dtype=torch.float32).cuda(),
            'right_img': torch.as_tensor(right_img, dtype=torch.float32).cuda(),
            'fu_mul_baseline': fu_mul_baseline.cuda(),
            'calibs_Proj': calibs_Proj.cuda(),
            'norm_coord_img': norm_coord_img.cuda(),
            # 'coord_img': coord_img.cuda(),
            'valid_mask_2d': valid_mask_2d.cuda(),
            'valids': valids.cuda()
        }
        batch_dict = self.model(batch_dict)
        # pred_bboxes = outs[0][0]['pred_boxes']
        # pred_scores = outs[0][0]['pred_scores']
        # pred_labels = outs[0][0]['pred_labels']
        # keys_to_keep = ['spatial_features_stride',
        #                     'spatial_features',
        #                     'spatial_features_2d',
        #                     'volume_features',
        #                     'batch_cls_preds',
        #                     'batch_box_preds']
        batch_cls_preds = batch_dict['lidar_outputs']['batch_cls_preds']
        batch_box_preds = batch_dict['lidar_outputs']['batch_box_preds']
        return batch_cls_preds, batch_box_preds


def register():
    """
    Register ONNX Runtime's built-in contrib ops.
	Should be run before torch.onnx.export().
	"""
    # @staticmethod
    # @parse_args('v', 'i', 'i', 'i', 'b')
    def grid_sampler(g, input, grid, mode, padding_mode, align_corners):
		# mode
		#   'bilinear'      : onnx::Constant[value={0}]
		#   'nearest'       : onnx::Constant[value={1}]
		#   'bicubic'       : onnx::Constant[value={2}]
		# padding_mode
		#   'zeros'         : onnx::Constant[value={0}]
		#   'border'        : onnx::Constant[value={1}]
		#   'reflection'    : onnx::Constant[value={2}]
        # grid = symbolic_helper._maybe_get_const(grid, "t")
        mode = symbolic_helper._maybe_get_const(mode, "i")
        padding_mode = symbolic_helper._maybe_get_const(padding_mode, "i")
        # mode_str = ["bilinear", "nearest", "bicubic"][mode]
        # padding_mode_str = ["zeros", "border", "reflection"][padding_mode]
        align_corners = int(symbolic_helper._maybe_get_const(align_corners, "b"))

		# From opset v13 onward, the output shape can be specified with
		# (N, C, H, W) (N, H_out, W_out, 2) => (N, C, H_out, W_out)
		# input_shape = input.type().sizes()
		# gird_shape = grid.type().sizes()
		# output_shape = input_shape[:2] + gird_shape[1:3]
		# g.op(...).setType(input.type().with_sizes(output_shape))
        inputs = [input, grid]
        kwargs = {
            "mode_i":mode, 
            "padding_mode_i":padding_mode, 
            "align_corners_i": align_corners}

        return g.op(
		    ## op name, modify here. not sure whether "com.microsoft::" is required
			"nvinfer1::GridSamplePluginDynamic",  
			*inputs, **kwargs
		)

    _reg(grid_sampler)

def _reg(symbolic_fn: typing.Callable):
	name = "::%s" % symbolic_fn.__name__
	torch.onnx.register_custom_op_symbolic(name, symbolic_fn, _OPSET_VERSION)
	_registered_ops.add(name)



def main():
    args, cfg = parse_config()
    if args.launcher == 'none':
        dist_test = False
        total_gpus = 1
    else:
        total_gpus, cfg.LOCAL_RANK = getattr(common_utils, 'init_dist_%s' % args.launcher)(
            args.tcp_port, args.local_rank, backend='nccl'
        )
        dist_test = True

    if args.batch_size is None:
        args.batch_size = cfg.OPTIMIZATION.BATCH_SIZE_PER_GPU
    else:
        assert args.batch_size % total_gpus == 0, 'Batch size should match the number of gpus'
        args.batch_size = args.batch_size // total_gpus

    if args.ckpt_id:
        assert args.exp_name

        output_dir = cfg.ROOT_DIR / 'outputs' / cfg.EXP_GROUP_PATH / (cfg.TAG + '.' + args.exp_name)
        args.ckpt = str(output_dir / 'ckpt' / 'checkpoint_epoch_{}.pth'.format(args.ckpt_id))
    elif args.ckpt:
        output_dir = Path(args.ckpt + ".eval")
        output_dir.mkdir(parents=True, exist_ok=True)
    else:
        raise ValueError("no ckpt specified")

    if not os.path.exists(output_dir):
        print('--- ! output_dir does not exist, please check the config file and the extra tag')
        output_dir.mkdir(parents=True, exist_ok=True)
    if args.eval_tag == 'eval_with_train':
        eval_output_dir = output_dir / 'eval' / 'eval_with_train'
    else:
        eval_output_dir = output_dir / 'eval' / args.eval_tag

    num_list = re.findall(r'\d+', args.ckpt) if args.ckpt is not None else []
    epoch_id = num_list[-1] if num_list.__len__() > 0 else 'no_number'
    eval_output_dir = eval_output_dir / ('epoch_%s' % epoch_id) / cfg.DATA_CONFIG.DATA_SPLIT['test']

    if not args.eval_tag :
        eval_output_dir = eval_output_dir / 'default'
    eval_output_dir.mkdir(parents=True, exist_ok=True)

    log_file = eval_output_dir / ('log_eval.txt')
    logger = common_utils.create_logger(log_file, rank=cfg.LOCAL_RANK)

    # log to file
    logger.info('**********************Start logging**********************')
    gpu_list = os.environ['CUDA_VISIBLE_DEVICES'] if 'CUDA_VISIBLE_DEVICES' in os.environ.keys() else 'ALL'
    logger.info('CUDA_VISIBLE_DEVICES={}'.format(gpu_list))
    logger.info('eval output dir: {}'.format(eval_output_dir))

    if dist_test:
        logger.info('total_batch_size: {}'.format(total_gpus * args.batch_size))
    for key, val in vars(args).items():
        logger.info('{:16} {}'.format(key, val))
    log_config_to_file(cfg, logger=logger)

    test_set, test_loader, sampler = build_dataloader(
        dataset_cfg=cfg.DATA_CONFIG,
        class_names=cfg.CLASS_NAMES,
        batch_size=args.batch_size,
        dist=dist_test, workers=args.workers, logger=logger, training=False
    )
    
    for i, batch_dict in enumerate(test_loader):
        if i == 0:
            a = batch_dict
        else:
            break



    model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=test_set)

    ## register for non-registered onnx ops -- grid sample
    register()

    with torch.no_grad():
        # eval_single_ckpt(model, test_loader, args, eval_output_dir, logger, epoch_id, dist_test=dist_test)
        model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=dist_test)
        model.cuda().eval()
        post_model = StereoDSGN_post(model=model)
        # batch_dict = {
        #     'left_img': torch.randn((1,3,480,928)),
        #     'right_img': torch.randn((1,3,480,928)),
        #     'calib': batch_dict['calib']
        # }
        N = 1
        left_img = torch.randn((1,3,480,928))
        right_img = torch.randn((1,3,480,928))
        calib_file = '/home/yaohan.lu/models/stereo-dsgn2/data/ww/training/calib/1692074207250000.txt'
        calib = calibration.Calibration(calib_file)
        fu_mul_baseline = torch.as_tensor([calib.fu_mul_baseline], dtype=left_img.dtype)
        calibs_Proj = torch.as_tensor([calib.P2 for i in range(N)], dtype=torch.float32)
        norm_coord_img = torch.tensor(np.load('norm_coord_img.npy')).unsqueeze(0)
        # coord_img = torch.tensor(np.load('coord_img.npy')).unsqueeze(0)
        valid_mask_2d = torch.tensor(np.load('valid_mask_2d.npy')).unsqueeze(0)
        valids = torch.tensor(np.load('valids.npy'))

        torch.onnx.export(
            post_model,
            (left_img, right_img, fu_mul_baseline, calibs_Proj, norm_coord_img, valid_mask_2d, valids),
            "dsgn_raw.onnx",
            export_params=True,
            opset_version=12,  # the ONNX version to export the model to
            do_constant_folding=True,  # whether to execute constant folding for optimization
            keep_initializers_as_inputs=True,
            input_names=['left_img', 'right_img', 'fu_mul_baseline', 'calibs_Proj', 'norm_coord_img', 'valid_mask_2d', 'valids'],  # the model's input names
            output_names = ['batch_cls_preds', 'batch_box_preds']  # the model's output names
        )
        print("convert sucessfully!")



if __name__ == '__main__':
    main()
